import os
import logging
try:
    import moxing as mox
except ModuleNotFoundError:
    logging.info("Moxing should be use of modelarts.")

from mindspore.train.callback._callback import Callback
from mindspore.train.serialization import load_checkpoint


def download(args, rank_id=0):
    """download ckpt from obs for finetune model."""
    obs_ckpt_path = args.finetune_ckpt
    ckpt_name = obs_ckpt_path.split("/")[-1]
    ckpt_path = args.ckpt_save_dir + str(rank_id)
    local_ckpt_path = os.path.join(ckpt_path, ckpt_name)
    mox.file.copy_parallel(obs_ckpt_path, local_ckpt_path)
    params_dict = None
    if os.path.exists(local_ckpt_path):
        args.logger.info(f"start loading {local_ckpt_path}.")
        params_dict = load_checkpoint(local_ckpt_path)
        args.logger.info(f"end loading {local_ckpt_path}.")
    return params_dict


class UploadCheckpoint(Callback):

    def __init__(self, src_dir, target_dir, upload_frequence=10):
        self.src_dir = src_dir
        self.target_dir = target_dir
        self.upload_frequence = upload_frequence

    def step_end(self, run_context):
        cb_params = run_context.original_args()
        if cb_params.cur_step_num % self.upload_frequence == 0 and os.listdir(self.src_dir):
            print("Starting upload model to obs!")
            print("Find ckpt file:", os.listdir(self.src_dir))
            print("Target dir is:{}".format(self.target_dir))
            try:
                mox.file.copy_parallel(src_url=self.src_dir,
                                       dst_url=self.target_dir)
                print("Upload ckpt succeed!")
            except Exception as e:
                print("Upload ckpt failed!", e)
